#!/bin/bash

# Ubuntu桌面环境 - Linux启动脚本
# 支持Ubuntu, CentOS, RHEL, Debian等Linux发行版

set -e  # 遇到错误立即退出

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 打印带颜色的消息
print_message() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

print_header() {
    echo -e "${BLUE}================================${NC}"
    echo -e "${BLUE}  Ubuntu桌面环境启动脚本${NC}"
    echo -e "${BLUE}================================${NC}"
}

# 检查是否为root用户
check_root() {
    if [[ $EUID -eq 0 ]]; then
        print_warning "检测到root用户，建议使用普通用户运行Docker"
        read -p "是否继续？(y/N): " -n 1 -r
        echo
        if [[ ! $REPLY =~ ^[Yy]$ ]]; then
            exit 1
        fi
    fi
}

# 检查Docker是否安装
check_docker() {
    print_message "检查Docker安装状态..."
    
    if ! command -v docker &> /dev/null; then
        print_error "Docker未安装，请先安装Docker"
        echo "安装命令参考："
        echo "  Ubuntu/Debian: sudo apt-get install docker.io docker-compose"
        echo "  CentOS/RHEL: sudo yum install docker docker-compose"
        echo "  Arch Linux: sudo pacman -S docker docker-compose"
        exit 1
    fi
    
    if ! command -v docker-compose &> /dev/null; then
        print_warning "docker-compose未安装，尝试安装..."
        if command -v apt-get &> /dev/null; then
            sudo apt-get install -y docker-compose
        elif command -v yum &> /dev/null; then
            sudo yum install -y docker-compose
        elif command -v pacman &> /dev/null; then
            sudo pacman -S docker-compose
        else
            print_error "无法自动安装docker-compose，请手动安装"
            exit 1
        fi
    fi
}

# 检查Docker服务状态
check_docker_service() {
    print_message "检查Docker服务状态..."
    
    if ! docker info &> /dev/null; then
        print_warning "Docker服务未运行，尝试启动..."
        
        # 尝试启动Docker服务
        if command -v systemctl &> /dev/null; then
            sudo systemctl start docker
            sudo systemctl enable docker
        elif command -v service &> /dev/null; then
            sudo service docker start
        else
            print_error "无法启动Docker服务，请手动启动"
            exit 1
        fi
        
        # 等待服务启动
        sleep 3
        
        # 再次检查
        if ! docker info &> /dev/null; then
            print_error "Docker服务启动失败"
            exit 1
        fi
    fi
    
    print_message "Docker服务运行正常"
}

# 检查网络冲突
check_network_conflicts() {
    print_message "检查网络配置..."
    
    local has_conflicts=false
    
    # 检查是否存在冲突的网络
    if docker network ls | grep -q "ubuntu-desktop_desktop-network"; then
        print_warning "检测到旧的网络配置，可能存在冲突"
        has_conflicts=true
    fi
    
    # 检查子网冲突
    local existing_subnets=($(docker network ls --format "{{.Name}}" | grep -E "(desktop|ubuntu)" | xargs -I {} docker network inspect {} --format "{{range .IPAM.Config}}{{.Subnet}}{{end}}" 2>/dev/null))
    
    for subnet in "${existing_subnets[@]}"; do
        if [[ "$subnet" == "172.20.0.0/16" ]] || [[ "$subnet" == "172.21.0.0/16" ]]; then
            print_warning "检测到子网冲突: $subnet"
            has_conflicts=true
        fi
    done
    
    if [ "$has_conflicts" = true ]; then
        print_warning "检测到网络冲突，建议清理后重新启动"
        read -p "是否自动清理网络冲突？(y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            cleanup_network_conflicts
        fi
    fi
}

# 清理网络冲突
cleanup_network_conflicts() {
    print_message "清理网络冲突..."
    
    # 停止相关容器
    docker-compose -f docker-compose.lightweight.yml down 2>/dev/null || true
    docker-compose down 2>/dev/null || true
    
    # 删除冲突的网络
    docker network rm ubuntu-desktop_desktop-network 2>/dev/null || true
    docker network rm ubuntu-desktop_desktop-network-lightweight 2>/dev/null || true
    docker network rm ubuntu-desktop_desktop-network-standard 2>/dev/null || true
    
    # 清理未使用的网络
    docker network prune -f
    
    print_message "网络冲突清理完成"
}

# 检查网络冲突
check_network_conflicts() {
    print_message "检查网络配置..."
    
    local has_conflicts=false
    
    # 检查是否存在冲突的网络
    if docker network ls | grep -q "ubuntu-desktop_desktop-network"; then
        print_warning "检测到旧的网络配置，可能存在冲突"
        has_conflicts=true
    fi
    
    # 检查子网冲突
    local existing_subnets=($(docker network ls --format "{{.Name}}" | grep -E "(desktop|ubuntu)" | xargs -I {} docker network inspect {} --format "{{range .IPAM.Config}}{{.Subnet}}{{end}}" 2>/dev/null))
    
    for subnet in "${existing_subnets[@]}"; do
        if [[ "$subnet" == "172.20.0.0/16" ]] || [[ "$subnet" == "172.21.0.0/16" ]]; then
            print_warning "检测到子网冲突: $subnet"
            has_conflicts=true
        fi
    done
    
    if [ "$has_conflicts" = true ]; then
        print_warning "检测到网络冲突，建议清理后重新启动"
        read -p "是否自动清理网络冲突？(y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            cleanup_network_conflicts
        fi
    fi
}

# 清理网络冲突
cleanup_network_conflicts() {
    print_message "清理网络冲突..."
    
    # 停止相关容器
    docker-compose -f docker-compose.lightweight.yml down 2>/dev/null || true
    docker-compose down 2>/dev/null || true
    
    # 删除冲突的网络
    docker network rm ubuntu-desktop_desktop-network 2>/dev/null || true
    docker network rm ubuntu-desktop_desktop-network-lightweight 2>/dev/null || true
    docker network rm ubuntu-desktop_desktop-network-standard 2>/dev/null || true
    
    # 清理未使用的网络
    docker network prune -f
    
    print_message "网络冲突清理完成"
}

# 检查系统资源
check_system_resources() {
    print_message "检查系统资源..."
    
    # 检查内存
    local mem_total=$(free -m | awk 'NR==2{printf "%.0f", $2}')
    if [ $mem_total -lt 2048 ]; then
        print_warning "系统内存不足2GB，可能影响性能"
    else
        print_message "系统内存: ${mem_total}MB"
    fi
    
    # 检查磁盘空间
    local disk_free=$(df -BG . | awk 'NR==2{print $4}' | sed 's/G//')
    if [ $disk_free -lt 8 ]; then
        print_warning "磁盘空间不足8GB，建议清理空间"
    else
        print_message "可用磁盘空间: ${disk_free}GB"
    fi
    
    # 检查CPU核心数
    local cpu_cores=$(nproc)
    print_message "CPU核心数: ${cpu_cores}"
}

# 清理Docker环境
cleanup_docker() {
    print_message "清理Docker环境..."
    
    # 清理未使用的镜像
    docker image prune -f
    
    # 清理未使用的容器
    docker container prune -f
    
    # 清理未使用的卷
    docker volume prune -f
    
    # 清理未使用的网络
    docker network prune -f
    
    # 清理构建缓存
    docker builder prune -f
    
    print_message "Docker环境清理完成"
}

# 选择版本
select_version() {
    echo
    echo "请选择要启动的版本："
    echo "1) 轻量级版本 (推荐，占用空间少，启动快)"
    echo "2) 标准版本 (功能完整，但占用空间多)"
    echo "3) 清理环境并退出"
    echo
    read -p "请输入选择 (1-3): " choice
    
    case $choice in
        1)
            VERSION="lightweight"
            COMPOSE_FILE="docker-compose.lightweight.yml"
            CONTAINER_NAME="ubuntu-desktop-lightweight"
            ;;
        2)
            VERSION="standard"
            COMPOSE_FILE="docker-compose.yml"
            CONTAINER_NAME="ubuntu-desktop-xfce"
            ;;
        3)
            cleanup_docker
            print_message "环境清理完成，退出"
            exit 0
            ;;
        *)
            print_error "无效选择，默认使用轻量级版本"
            VERSION="lightweight"
            COMPOSE_FILE="docker-compose.lightweight.yml"
            CONTAINER_NAME="ubuntu-desktop-lightweight"
            ;;
    esac
    
    print_message "选择版本: ${VERSION}"
}

# 构建镜像
build_image() {
    print_message "构建Docker镜像..."
    
    if [ "$VERSION" = "lightweight" ]; then
        docker-compose -f $COMPOSE_FILE build --no-cache
    else
        docker-compose build --no-cache
    fi
    
    if [ $? -ne 0 ]; then
        print_error "镜像构建失败"
        exit 1
    fi
    
    print_message "镜像构建成功"
}

# 启动容器
start_container() {
    print_message "启动容器..."
    
    if [ "$VERSION" = "lightweight" ]; then
        docker-compose -f $COMPOSE_FILE up -d
    else
        docker-compose up -d
    fi
    
    if [ $? -ne 0 ]; then
        print_error "容器启动失败"
        exit 1
    fi
    
    print_message "容器启动成功"
}

# 等待容器就绪
wait_for_container() {
    print_message "等待容器就绪..."
    
    local max_attempts=30
    local attempt=1
    
    while [ $attempt -le $max_attempts ]; do
        if docker ps | grep -q $CONTAINER_NAME; then
            local status=$(docker inspect --format='{{.State.Status}}' $CONTAINER_NAME 2>/dev/null)
            if [ "$status" = "running" ]; then
                print_message "容器运行正常"
                return 0
            fi
        fi
        
        echo -n "."
        sleep 2
        attempt=$((attempt + 1))
    done
    
    print_error "容器启动超时"
    return 1
}

# 显示访问信息
show_access_info() {
    echo
    echo -e "${GREEN}================================${NC}"
    echo -e "${GREEN}  Ubuntu桌面环境启动成功！${NC}"
    echo -e "${GREEN}================================${NC}"
    echo
    
    if [ "$VERSION" = "lightweight" ]; then
        echo "轻量级版本访问信息："
        echo "  VNC访问: localhost:5901"
        echo "  SSH访问: ssh -p 2223 ubuntu@localhost"
    else
        echo "标准版本访问信息："
        echo "  VNC访问: localhost:5900"
        echo "  SSH访问: ssh -p 2222 ubuntu@localhost"
    fi
    
    echo
    echo "本地X11显示（Linux特有）："
    echo "  1. 允许X11连接: xhost +local:docker"
    echo "  2. 启动桌面: docker exec -it $CONTAINER_NAME startxfce4"
    echo
    echo "管理命令："
    echo "  查看状态: docker-compose -f $COMPOSE_FILE ps"
    echo "  查看日志: docker-compose -f $COMPOSE_FILE logs -f"
    echo "  停止服务: docker-compose -f $COMPOSE_FILE down"
    echo "  进入容器: docker exec -it $CONTAINER_NAME bash"
    echo
}

# 主函数
main() {
    print_header
    
    # 检查系统环境
    check_root
    check_docker
    check_docker_service
    
    # 检查网络冲突
    check_network_conflicts
    
    # 检查系统资源
    check_system_resources
    
    # 选择版本
    select_version
    
    # 清理环境（可选）
    read -p "是否清理Docker环境？(y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        cleanup_docker
    fi
    
    # 构建和启动
    build_image
    start_container
    
    # 等待容器就绪
    if wait_for_container; then
        show_access_info
    else
        print_error "启动失败，请检查日志"
        if [ "$VERSION" = "lightweight" ]; then
            docker-compose -f $COMPOSE_FILE logs
        else
            docker-compose logs
        fi
        exit 1
    fi
}

# 捕获信号
trap 'print_error "脚本被中断"; exit 1' INT TERM

# 运行主函数
main "$@" 